import cv2
import numpy as np
from PIL import Image
from skimage.segmentation import slic
from skimage.color import rgb2lab
from collections import Counter
import requests
from io import BytesIO


# 从URL加载图片
def load_image_from_url(url):
    """
    从URL加载图片并转换为RGB格式
    """
    try:
        response = requests.get(url)
        img = Image.open(BytesIO(response.content)).convert("RGB")
        img_rgb = np.array(img)
        return img_rgb
    except Exception as e:
        print(f"Error loading image from URL: {e}")
        return None


# 1. 图片预分类（贴图、产品图、效果图）
def image_pre_classification(image_rgb):
    """
    图片预分类：将图片分为贴图、产品图、效果图
    使用简单的特征（如边缘密度、颜色分布）进行分类
    """
    # 转换为BGR格式以使用OpenCV
    img_bgr = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2BGR)

    # 边缘检测：贴图通常边缘较少，产品图边缘清晰，效果图边缘复杂
    gray = cv2.cvtColor(img_bgr, cv2.COLOR_BGR2GRAY)
    edges = cv2.Canny(gray, 100, 200)
    edge_density = np.sum(edges) / (edges.shape[0] * edges.shape[1])

    # 颜色直方图：贴图颜色单一，产品图中等，效果图复杂
    hist = cv2.calcHist([img_bgr], [0, 1, 2], None, [8, 8, 8], [0, 256, 0, 256, 0, 256])
    hist_norm = hist / hist.sum()
    color_complexity = -np.sum(hist_norm * np.log2(hist_norm + 1e-10))

    # 分类逻辑（阈值需根据实际数据调整）
    if edge_density < 0.05 and color_complexity < 2:
        return "贴图", img_bgr
    elif edge_density < 0.2 and color_complexity < 5:
        return "产品图", img_bgr
    else:
        return "效果图", img_bgr


# 2. 图片主体识别与裁剪（CPU实现）
def subject_detection(image_rgb, target_category):
    """
    图片主体识别：根据传入的分类字段提取特定类别的目标主体并裁剪（CPU实现）
    使用skimage的SLIC分割和OpenCV的Canny边缘检测
    """
    # SLIC分割（超像素分割）
    segments = slic(image_rgb, n_segments=100, compactness=10, sigma=1)

    # 边缘检测
    gray = cv2.cvtColor(image_rgb, cv2.COLOR_RGB2GRAY)
    edges = cv2.Canny(gray, 100, 200)

    # 根据传入的分类字段选择主体区域
    if target_category == "贴图":
        # 贴图：主体通常居中，背景单一
        center_segment = segments[image_rgb.shape[0] // 2, image_rgb.shape[1] // 2]
        mask = (segments == center_segment).astype(np.uint8) * 255
    elif target_category == "产品图":
        # 产品图：主体边缘清晰，找最大边缘区域
        contours, _ = cv2.findContours(edges, cv2.RETR_EXTERNAL, cv2.CHAIN_APPROX_SIMPLE)
        if contours:
            largest_contour = max(contours, key=cv2.contourArea)
            mask = np.zeros_like(gray)
            cv2.drawContours(mask, [largest_contour], -1, 255, thickness=cv2.FILLED)
        else:
            mask = np.ones_like(gray) * 255
    elif target_category == "效果图":
        # 效果图：多主体，找显著区域（简化处理）
        mask = (edges > 0).astype(np.uint8) * 255
    else:
        raise ValueError(f"Unsupported target category: {target_category}")

    # 找到主体的边界框
    coords = np.column_stack(np.where(mask > 0))
    if len(coords) == 0:
        return image_rgb  # 如果没找到主体，返回原图
    y_min, x_min = coords.min(axis=0)
    y_max, x_max = coords.max(axis=0)

    # 裁剪主体
    cropped_image = image_rgb[y_min:y_max + 1, x_min:x_max + 1]
    return cropped_image


# 3. 颜色提取（CPU实现）
def color_extraction(image_rgb):
    """
    颜色提取：使用K-means聚类提取主要颜色（CPU实现）
    返回RGB值
    """
    # 转换为LAB颜色空间，便于聚类
    image_lab = rgb2lab(image_rgb)
    pixels = image_lab.reshape(-1, 3)

    # K-means聚类（CPU实现）
    def kmeans_cpu(data, k, max_iter=10):
        # 随机初始化质心
        centroids = data[np.random.choice(len(data), k, replace=False)]
        for _ in range(max_iter):
            # 计算距离
            distances = np.sum((data[:, None] - centroids[None, :]) ** 2, axis=2)
            labels = np.argmin(distances, axis=1)
            # 更新质心
            for i in range(k):
                if np.sum(labels == i) > 0:
                    centroids[i] = np.mean(data[labels == i], axis=0)
        return labels, centroids

    # 提取主要颜色（假设提取3种主要颜色）
    labels, centroids = kmeans_cpu(pixels, k=3)

    # 转换回RGB
    from skimage.color import lab2rgb
    centroids_rgb = lab2rgb(centroids.reshape(-1, 1, 3)).reshape(-1, 3) * 255
    return centroids_rgb.astype(int)


# 主流程
def handle(image_url, target_category):
    """
    主流程：从URL加载图片，提取指定类别的目标主体，并提取颜色
    :param image_url: 图片URL地址
    :param target_category: 目标分类字段（贴图、产品图、效果图）
    """
    # 1. 从URL加载图片
    img_rgb = load_image_from_url(image_url)
    if img_rgb is None:
        print("Failed to load image. Exiting.")
        return

    # 2. 图片预分类（仅用于记录，不影响主体提取）
    pre_category, _ = image_pre_classification(img_rgb)
    print(f"预分类结果: {pre_category}")

    # 3. 图片主体识别与裁剪（使用传入的分类字段）
    cropped_img = subject_detection(img_rgb, pre_category)
    print(f"主体裁剪完成，尺寸: {cropped_img.shape}")

    # 4. 颜色提取
    main_colors = color_extraction(cropped_img)
    print("主要颜色 (RGB):")
    for i, color in enumerate(main_colors):
        print(f"颜色 {i + 1}: {color}")
